import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# 数据预处理
# 创建一个包含多个数据预处理步骤的组合
transform = transforms.Compose([
    transforms.ToTensor(),  # 将图像转换为张量
    transforms.Normalize((0.1307,), (0.3081,))  # 对张量进行归一化
])

# 加载MNIST数据集
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)

# 划分训练集和验证集
train_size = int(0.8 * len(train_dataset))
val_size = len(train_dataset) - train_size
train_dataset, val_dataset = torch.utils.data.random_split(train_dataset, [train_size, val_size])

# 创建数据加载器
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

# 定义模型
class MLP(nn.Module):
    def __init__(self):
        super(MLP, self).__init__()
        self.flatten = nn.Flatten()  # 创建一个将张量转换为一维张量的层
        self.fc1 = nn.Linear(784, 20)  # 创建一个全连接层，输入大小为784，输出大小为20
        self.relu = nn.ReLU()  # 创建一个ReLU激活函数层
        self.fc2 = nn.Linear(20, 10)  # 创建一个全连接层，输入大小为20，输出大小为10

    def forward(self, x):
        x = self.flatten(x)  # 将张量转换为一维张量
        x = self.fc1(x)  # 通过全连接层
        x = self.relu(x)  # 通过ReLU激活函数
        x = self.fc2(x)  # 通过全连接层
        return x

# 创建模型实例
model = MLP()

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()  # 创建一个交叉熵损失函数
optimizer = optim.Adam(model.parameters(), lr=0.001)  # 创建一个Adam优化器

# 训练模型
num_epochs = 20  # 设置训练轮数
for epoch in range(num_epochs):
    model.train()  # 将模型设置为训练模式
    train_loss = 0.0  # 初始化训练损失
    train_acc = 0.0  # 初始化训练准确率
    for images, labels in train_loader:  # 遍历训练数据加载器
        outputs = model(images)  # 通过模型计算输出
        loss = criterion(outputs, labels)  # 计算损失
        optimizer.zero_grad()  # 清除梯度
        loss.backward()  # 反向传播
        optimizer.step()  # 更新参数
        train_loss += loss.item() * images.size(0)  # 累加训练损失
        _, predicted = torch.max(outputs.data, 1)  # 获取最大值所在的位置
        train_acc += (predicted == labels).sum().item()  # 累加训练准确率
    train_loss /= len(train_dataset)  # 计算平均训练损失
    train_acc /= len(train_dataset)  # 计算平均训练准确率

    model.eval()  # 将模型设置为评估模式
    val_loss = 0.0  # 初始化验证损失
    val_acc = 0.0  # 初始化验证准确率
    with torch.no_grad():  # 不计算梯度

        for images, labels in val_loader:  # 遍历验证数据加载器
            outputs = model(images)  # 通过模型计算输出
            loss = criterion(outputs, labels)  # 计算损失
            val_loss += loss.item() * images.size(0)  # 累加验证损失
            _, predicted = torch.max(outputs.data, 1)  # 获取最大值所在的位置
            val_acc += (predicted == labels).sum().item()  # 累加验证准确率
    val_loss /= len(val_dataset)  # 计算平均验证损失
    val_acc /= len(val_dataset)  # 计算平均验证准确率

    print(f"Epoch [{epoch + 1}/{num_epochs}], Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.4f}")
    print(f"Epoch [{epoch + 1}/{num_epochs}], Val Loss: {val_loss:.4f}, Val Acc: {val_acc:.4f}")

# 测试模型
model.eval()  # 将模型设置为评估模式
test_loss = 0.0  # 初始化测试损失
test_acc = 0.0  # 初始化测试准确率
with torch.no_grad():  # 不计算梯度
    for images, labels in test_loader:  # 遍历测试数据加载器
        outputs = model(images)  # 通过模型计算输出
        loss = criterion(outputs, labels)  # 计算损失
        test_loss += loss.item() * images.size(0)  # 累加测试损失
        _, predicted = torch.max(outputs.data, 1)  # 获取最大值所在的位置
        test_acc += (predicted == labels).sum().item()  # 累加测试准确率
test_loss /= len(test_dataset)  # 计算平均测试损失
test_acc /= len(test_dataset)  # 计算平均测试准确率

print(f"Test Loss: {test_loss:.4f}, Test Acc: {test_acc:.4f}")
